#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelFastGeluGrad
{
public:
  __aicore__ inline KernelFastGeluGrad() {}
  __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR z, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t typeKey)
  {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->typeKey = typeKey;
    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;
    dyGm.SetGlobalBuffer((__gm__ DTYPE_X *)dy,
                        this->blockLength);
    xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x,
                        this->blockLength);
    zGm.SetGlobalBuffer((__gm__ DTYPE_X *)z,
                          this->blockLength);

    isbig = 0;
    pipe.InitBuffer(inQueueDY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(calcBuf, this->tileLength * sizeof(DTYPE_X));
  }
  __aicore__ inline void Process()
  {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++)
    {
      if(BUFFER_NUM == 2 && lasttileLength <= tileLength && i == loopCount - 2)continue;
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
  }

private:
  __aicore__ inline void CopyIn(int32_t progress)
  {
    LocalTensor<DTYPE_X> dyLocal = inQueueDY.AllocTensor<DTYPE_X>();
    LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();

    if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
        (progress == (this->tileNum * BUFFER_NUM - 1))) {
      DataCopy(
          dyLocal[0],
          dyGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
          (this->tileLength));
      DataCopy(
          xLocal[0],
          xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
          (this->tileLength));

    }
    else {
      DataCopy(dyLocal[0], dyGm[progress * (this->tileLength)],
                (this->tileLength));
      DataCopy(xLocal[0], xGm[progress * (this->tileLength)],
                (this->tileLength));
    }
    inQueueDY.EnQue(dyLocal);
    inQueueX.EnQue(xLocal);
    
  }
  __aicore__ inline void Compute(int32_t progress)
  {
    LocalTensor<DTYPE_X> dyLocal = inQueueDY.DeQue<DTYPE_X>();
    LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
    LocalTensor<DTYPE_X> zLocal = outQueueZ.AllocTensor<DTYPE_X>();
    LocalTensor<DTYPE_X> tmpLocal = calcBuf.Get<DTYPE_X>();
    if(typeKey == 0){
      if(isbig == 0){
        ReduceMax(zLocal, xLocal, tmpLocal, this->tileLength, false);
        if((float)zLocal.GetValue(0) > (float)20.0){
          isbig = 1;
        }else{
          isbig = -1;
        }
      }
      if(isbig == 1){        
        Abs(zLocal, xLocal, this->tileLength);
        Sub(tmpLocal, xLocal, zLocal, this->tileLength);
        Muls(tmpLocal, tmpLocal, (DTYPE_X)1.702, this->tileLength);
        Exp(tmpLocal, tmpLocal, this->tileLength);
        Muls(zLocal, zLocal, (DTYPE_X)-1.702, this->tileLength);
        Exp(zLocal, zLocal, this->tileLength);
        Mul(xLocal, xLocal, zLocal, this->tileLength);
        Muls(xLocal, xLocal, (DTYPE_X)1.702, this->tileLength);
        Add(xLocal, zLocal, xLocal, this->tileLength);
        Add(tmpLocal, tmpLocal, xLocal, this->tileLength);
        Adds(xLocal, zLocal, (DTYPE_X)1.0, this->tileLength);
        Mul(xLocal, xLocal, xLocal, this->tileLength);
        Div(zLocal, tmpLocal, xLocal, this->tileLength);
        Mul(zLocal, dyLocal, zLocal, this->tileLength);
      }else{
        Muls(xLocal, xLocal, (DTYPE_X)-1.702, this->tileLength); // x = -1.702x
        Exp(tmpLocal, xLocal, this->tileLength); // tmp = e^(-1.702x)
        Mul(zLocal, tmpLocal, xLocal, this->tileLength); // z = -1.702x e^(-1.702x)
        Adds(tmpLocal, tmpLocal, (DTYPE_X)1.0, this->tileLength); // tmp = e^(-1.702x)+1
        Muls(zLocal, zLocal, (DTYPE_X)-1.0, this->tileLength);
        Add(zLocal, tmpLocal, zLocal, this->tileLength); // z = e^(-1.702x)+1+1.702xe^(-1.702x)
        Mul(tmpLocal, tmpLocal, tmpLocal, this->tileLength); // (e^(-1.702x)+1)^2
        Div(zLocal, zLocal, tmpLocal, this->tileLength);
        Mul(zLocal, dyLocal, zLocal, this->tileLength);
      }
    }else{
      Abs(zLocal, xLocal, this->tileLength);
      Sub(tmpLocal, xLocal, zLocal, this->tileLength);
      Muls(tmpLocal, tmpLocal, (DTYPE_X)1.702, this->tileLength);
      Exp(tmpLocal, tmpLocal, this->tileLength);
      Muls(zLocal, zLocal, (DTYPE_X)-1.702, this->tileLength);
      Exp(zLocal, zLocal, this->tileLength);
      Mul(xLocal, xLocal, zLocal, this->tileLength);
      Muls(xLocal, xLocal, (DTYPE_X)1.702, this->tileLength);
      Add(xLocal, zLocal, xLocal, this->tileLength);
      Add(tmpLocal, tmpLocal, xLocal, this->tileLength);
      Adds(xLocal, zLocal, (DTYPE_X)1.0, this->tileLength);
      Mul(xLocal, xLocal, xLocal, this->tileLength);
      Div(zLocal, tmpLocal, xLocal, this->tileLength);
      Mul(zLocal, dyLocal, zLocal, this->tileLength);
    }
    
    outQueueZ.EnQue<DTYPE_X>(zLocal);
    inQueueDY.FreeTensor(dyLocal);
    inQueueX.FreeTensor(xLocal);
  }
  __aicore__ inline void CopyOut(int32_t progress)
  {
    LocalTensor<DTYPE_X> zLocal = outQueueZ.DeQue<DTYPE_X>();

    if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
        (progress == (this->tileNum * BUFFER_NUM - 1))) {
      DataCopy(
          zGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
          zLocal, (this->tileLength));
    }

    else {
      DataCopy(zGm[progress * (this->tileLength)], zLocal,
                (this->tileLength));
    }

    outQueueZ.FreeTensor(zLocal);
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueDY;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
  TBuf<TPosition::VECCALC> calcBuf;
  GlobalTensor<DTYPE_X> dyGm;
  GlobalTensor<DTYPE_X> xGm;
  GlobalTensor<DTYPE_Y> zGm;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t typeKey;
  uint32_t isbig;
};
extern "C" __global__ __aicore__ void fast_gelu_grad(GM_ADDR dy, GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelFastGeluGrad op;

    op.Init(dy, x, z, tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, tiling_data.typeKey);
    op.Process();
    // TODO: user kernel impl
}
#ifndef __CCE_KT_TEST__
void fast_gelu_grad_do(uint32_t blockDim, void* l2ctrl, void* stream,
                       uint8_t* dy, uint8_t* x, uint8_t* z,
                       uint8_t* workspace, uint8_t* tiling) {
  fast_gelu_grad<<<blockDim, l2ctrl, stream>>>(dy, x, z, workspace, tiling);
}
#endif